'''
Created on Jan 30, 2011

@author: Yevgeniy Gindin

Contains modified code from elasticwulf (http://code.google.com/p/elasticwulf) 
Copyright (c) 2007 DataWrangling. All rights reserved.

'''
import sys
import getopt
import ec2
import os
import commands
import socket
from config import *

help_message = '''
SYNOPSIS
   ec2-mpi-config (ec2-MPI-configuration)
   ec2-mpi-config [-rsh False] 
GENERAL NOTES
   #
DESCRIPTION
    Before running, change the configuration strings in EC2config.py to match your AWS keys

    This script will:

    1. Generate list of your running or pending Elasticwulf node instances on EC2
    2. Select the first node as the master, and scp over the host files to all machines in the MPI cluster
    3. If the master node is still in a "pending" state, the script will sleep and try again until cancelled by the user


    Note: *The ssh variable StrictHostKeyChecking is set to "no". This is an evil hack to avoid the tedious adding of each compute
    node host when managing a large cluster...I'm assuming these EC2 nodes will only connect to eachother and S3, please be careful.

    See http://www.securityfocus.com/infocus/1806
OPTIONS
    #
'''

class Usage(Exception):
    def __init__(self, msg):
        self.msg = msg

def main(argv=None):
    if argv is None:
        argv = sys.argv
    try:
        try:
            opts, args = getopt.getopt(argv[1:], "ho:v", ["help", "output="])
        except getopt.error, msg:
            raise Usage(msg)
    
        # option processing
        for option, value in opts:
            if option == "-v":
                verbose = True
            if option in ("-h", "--help"):
                raise Usage(help_message)
            if option in ("-o", "--output"):
                output = value
                
        configure()
    
    except Usage, err:
        print >> sys.stderr, sys.argv[0].split("/")[-1] + ": " + str(err.msg)
        print >> sys.stderr, "\t for help use --help"
        return 2


def remote(host, cmd='scp',
           user='root', src=None, dest=None,
           credential=None,
           silent=False, test=False):
    d = {
         'cmd':cmd,
         'user':user,
         'src':src,
         'dest':dest,
         'host':host,
         }

    d['switches'] = ''
    if credential:
            d['switches'] = '-i %s' % credential
    
    if cmd == 'scp':
            
        template = '%(cmd)s %(switches)s -o "StrictHostKeyChecking no" %(src)s %(user)s@%(host)s:%(dest)s' 
    else:
        template = 'ssh -o "StrictHostKeyChecking no" %(user)s@%(host)s "%(cmd)s" '

    cmdline = template % d  

    if not silent:
        print "\n",cmdline,"\n\n"
    if not test:
        os.system(cmdline)



def configure():
    import shutil
    conn = ec2.AWSAuthConnection(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY)
    instance_response=conn.describe_instances()
    parsed_response=instance_response.parse()  
    mpi_instances=[]
    mpi_amis=[]
    mpi_hostnames=[]
    mpi_externalnames=[]
    machine_state=[]

    for chunk in parsed_response:
        if chunk[0]=='INSTANCE':
            if chunk[-1]=='running' or chunk[-1]=='pending':
                if (chunk[2] == IMAGE_ID):
                    mpi_instances.append(chunk[1])
                    mpi_amis.append(chunk[2])
                    mpi_externalnames.append(chunk[3])  
                    mpi_hostnames.append(chunk[4])                          
                    machine_state.append(chunk[-1])

    # list the nodes
    if len(mpi_instances) > 0:
        print "\n---- MPI Cluster Details ----"
        print "Number of nodes = %s" % len(mpi_instances)

        for i, machine in enumerate(mpi_instances):
            print "Instance= %s external_name = %s hostname= %s state= %s" %(machine, mpi_externalnames[i], mpi_hostnames[i], machine_state[i])

    # Configure mpd.hosts: write out the hostnames to a mpd.hosts file
    print "\nWriting out mpd.hosts file"
    hosts_file= "mpd.hosts"
    output=open(hosts_file,'w')
    for host in mpi_hostnames:
        print >> output, "%s" % host
    output.close()
    
    homedir = os.path.expanduser( '~' )
    try:
        # check if key already in id_rsa...
        rsakeys = open(homedir + "/.ssh/id_rsa", 'r').read()
        keyfile = open(KEY_LOCATION,'r').read()
        keyval = keyfile.split('-----')[2]
        key_index = rsakeys.split('-----').index(keyval)
        print "Key already in id_rsa..."
    except ValueError:
        # if not, add it...
        print "Appending key to id_rsa..."
        os.system("cat %s >> %s/.ssh/id_rsa" % (KEY_LOCATION, homedir) )
        os.system("chmod 600 %s/.ssh/id_rsa" % homedir)

    for host in mpi_externalnames:
        remote(host, src=hosts_file, dest="/root/")

if __name__ == "__main__":
    sys.exit(main())
